import pandas as pd
import os
import numpy as np
from datetime import datetime
import plotly.graph_objects as go
from plotly.subplots import make_subplots


# --- 数据加载和处理函数 (与您的版本相同，无需修改) ---
def load_data_in_range(start_time_str: str, end_time_str: str, base_path: str, symbol: str,
                       timeframe: str) -> pd.DataFrame:
    """从指定的CSV文件中加载时间范围内的数据。"""
    start_date = pd.to_datetime(start_time_str)
    end_date = pd.to_datetime(end_time_str)
    data_folder_path = os.path.join(base_path, symbol, timeframe)
    if not os.path.isdir(data_folder_path):
        print(f"错误: 文件夹不存在 -> {data_folder_path}")
        return pd.DataFrame()

    date_range = pd.date_range(start=start_date.date(), end=end_date.date(), freq='D').strftime('%Y-%m').unique()
    if len(date_range) == 0: date_range = [start_date.strftime('%Y-%m')]
    all_dfs = []
    print(f"将在以下月份文件中查找数据: {list(date_range)}")
    for year_month in date_range:
        file_path = os.path.join(data_folder_path, f"{symbol.upper()}-{timeframe}-{year_month}.csv")
        if os.path.exists(file_path):
            print(f"正在读取文件: {file_path}")
            try:
                df = pd.read_csv(file_path, usecols=['open_time', 'open', 'high', 'low', 'close', 'volume'])
                all_dfs.append(df)
            except Exception as e:
                print(f"读取或处理文件 {file_path} 时出错: {e}")
        else:
            print(f"警告: 未找到文件 {file_path}")

    if not all_dfs:
        print("错误: 在指定时间范围内未找到任何数据文件。")
        return pd.DataFrame()

    combined_df = pd.concat(all_dfs, ignore_index=True)
    combined_df.drop_duplicates(subset=['open_time'], inplace=True)
    return combined_df


def process_and_filter_data(df: pd.DataFrame, start_time_str: str, end_time_str: str) -> pd.DataFrame:
    """数据预处理：转换时间、排序、筛选。"""
    if df.empty:
        return df
    df['open_time'] = pd.to_datetime(df['open_time'], unit='ms')
    df.set_index('open_time', inplace=True)
    numeric_cols = ['open', 'high', 'low', 'close', 'volume']
    df[numeric_cols] = df[numeric_cols].apply(pd.to_numeric, errors='coerce')
    df.dropna(subset=numeric_cols, inplace=True)
    df.sort_index(inplace=True)
    df_filtered = df.loc[start_time_str:end_time_str]
    print(f"数据处理完成。筛选出 {len(df_filtered)} 条K线。")
    return df_filtered


# --- 新增：核心功能函数 ---

def calculate_atr(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
    """计算ATR指标并添加到DataFrame中。"""
    df_copy = df.copy()
    df_copy['high_low'] = df_copy['high'] - df_copy['low']
    df_copy['high_close'] = np.abs(df_copy['high'] - df_copy['close'].shift())
    df_copy['low_close'] = np.abs(df_copy['low'] - df_copy['close'].shift())
    # TR (True Range) 是三者中的最大值
    df_copy['tr'] = df_copy[['high_low', 'high_close', 'low_close']].max(axis=1)
    # ATR 是 TR 的移动平均
    df_copy['atr'] = df_copy['tr'].rolling(window=period).mean()
    print(f"已计算周期为 {period} 的ATR。")
    return df_copy.drop(columns=['high_low', 'high_close', 'low_close', 'tr'])


def find_and_merge_boxes(df: pd.DataFrame, n: int, k: float, slope_threshold: float) -> list:
    """
    根据您的量化思路查找并合并箱体。
    :param df: 包含K线和ATR数据的数据帧
    :param n: 识别箱体的K线周期 (window size)
    :param k: ATR倍数阈值
    :param slope_threshold: 标准化后的斜率绝对值阈值
    :return: 一个包含合并后箱体坐标的列表
    """
    if 'atr' not in df.columns or df['atr'].isnull().all():
        print("错误: DataFrame中缺少有效的'atr'列。无法识别箱体。")
        return []

    print(f"\n--- 开始识别箱体 (周期={n}, ATR倍数<{k}, 标准化斜率<{slope_threshold}) ---")
    potential_boxes = []
    # 1. 遍历数据，找出所有满足条件的潜在箱体窗口
    for i in range(len(df) - n + 1):
        window = df.iloc[i:i + n]
        if window['atr'].isnull().any():
            continue

        max_high = window['high'].max()
        min_low = window['low'].min()
        box_height = max_high - min_low
        # 使用窗口内最后一个ATR值作为参考
        atr_val = window['atr'].iloc[-1]

        # 条件1: 最大值-最小值 < k * ATR
        if box_height < k * atr_val and box_height > 0:
            # 条件2: 检查箱体是否大致水平
            # 对窗口内的收盘价进行线性回归
            x = np.arange(n)
            y = window['close'].values
            slope, _ = np.polyfit(x, y, 1)

            # 斜率标准化：将斜率转换为“每根K线价格变化的百分比”
            # 这样可以使其不受币价本身高低的影响
            normalized_slope = slope / window['close'].mean()

            if abs(normalized_slope) < slope_threshold:
                potential_boxes.append({
                    'x0': window.index[0],
                    'y0': min_low,
                    'x1': window.index[-1],
                    'y1': max_high
                })

    print(f"找到了 {len(potential_boxes)} 个潜在的重叠箱体，开始合并...")

    # 2. 合并重叠或连续的箱体
    if not potential_boxes:
        print("未找到符合条件的箱体。")
        return []

    # 按开始时间排序
    potential_boxes.sort(key=lambda b: b['x0'])

    merged_boxes = [potential_boxes[0]]
    for current_box in potential_boxes[1:]:
        last_merged_box = merged_boxes[-1]
        # 如果当前箱体的开始时间在最后合并的箱体结束之前，说明它们重叠了
        if current_box['x0'] <= last_merged_box['x1']:
            # 合并两个箱体：结束时间取最晚的，价格范围取并集
            last_merged_box['x1'] = max(last_merged_box['x1'], current_box['x1'])
            last_merged_box['y0'] = min(last_merged_box['y0'], current_box['y0'])
            last_merged_box['y1'] = max(last_merged_box['y1'], current_box['y1'])
        else:
            # 如果不重叠，则将其作为一个新的箱体添加
            merged_boxes.append(current_box)

    print(f"合并完成，最终识别出 {len(merged_boxes)} 个独立箱体。")
    return merged_boxes


# --- 新增：专注绘制K线和箱体的绘图函数 ---
def plot_kline_with_boxes(df: pd.DataFrame, boxes: list, chart_title: str):
    """
    使用Plotly绘制K线图，并在图上绘制识别出的箱体。
    """
    if df.empty:
        print("数据为空，无法绘制图表。")
        return

    fig = go.Figure()

    # 1. 添加K线图
    fig.add_trace(go.Candlestick(x=df.index,
                                 open=df['open'], high=df['high'],
                                 low=df['low'], close=df['close'],
                                 name='OHLC'))

    # 2. 添加箱体形状
    for box in boxes:
        fig.add_shape(
            type="rect",
            x0=box['x0'], y0=box['y0'], x1=box['x1'], y1=box['y1'],
            line=dict(color="rgba(255, 165, 0, 0.8)", width=1), # 橙色边框
            fillcolor="rgba(255, 165, 0, 0.2)", # 橙色半透明填充
        )

    # 3. 更新布局
    start_str = df.index[0].strftime('%Y-%m-%d %H:%M')
    end_str = df.index[-1].strftime('%Y-%m-%d %H:%M')
    full_title = f'{chart_title} - 箱体识别 ({start_str} to {end_str})'

    fig.update_layout(
        title_text=full_title,
        xaxis_title="时间",
        yaxis_title="价格 (USDT)",
        xaxis_rangeslider_visible=False, # 隐藏底部范围滑块
        template="plotly_dark" # 使用深色主题
    )

    fig.show()


if __name__ == '__main__':
    # ==================== 用户配置区 ====================
    BASE_DATA_PATH = "F:/personal/binance_klines"
    SYMBOL = "ETHUSDT"
    TIMEFRAME = "1h"
    START_TIME = "2025-06-01 00:00:00"
    END_TIME = "2025-07-01 00:00:00"

    # --- 箱体识别参数 ---
    BOX_FINDING_PARAMS = {
        "n": 40,                 # 识别箱体的K线周期数 (窗口大小)
        "k": 6,                # 箱体高度必须小于 k * ATR
        "slope_threshold": 0.0001, # 标准化后的斜率绝对值阈值 (一个非常小的值代表接近水平)
        "atr_period": 40         # 计算ATR时使用的周期
    }
    # ====================================================

    print("--- K线与箱体识别程序启动 ---")
    # 1. 加载和处理数据
    raw_data_df = load_data_in_range(START_TIME, END_TIME, BASE_DATA_PATH, SYMBOL, TIMEFRAME)
    final_df = process_and_filter_data(raw_data_df, START_TIME, END_TIME)

    if not final_df.empty:
        # 2. 计算ATR指标
        final_df_with_atr = calculate_atr(final_df, period=BOX_FINDING_PARAMS["atr_period"])

        # 3. 识别并合并箱体
        identified_boxes = find_and_merge_boxes(
            final_df_with_atr,
            n=BOX_FINDING_PARAMS["n"],
            k=BOX_FINDING_PARAMS["k"],
            slope_threshold=BOX_FINDING_PARAMS["slope_threshold"]
        )

        # 4. 绘图
        chart_title_str = f"{SYMBOL} - {TIMEFRAME}"
        plot_kline_with_boxes(final_df_with_atr, identified_boxes, chart_title=chart_title_str)
        print("\n--- 图表生成完毕 ---")
    else:
        print("\n未能加载或处理任何数据，程序退出。请检查您的配置和文件路径。")